import matplotlib.pyplot as plt


decision_node = dict(boxstyle="sawtooth", fc="0.8")     # 决策点形状
leaf_node = dict(boxstyle="round4", fc="0.8")           # 叶子点形状
arrow_args = dict(arrowstyle="<-")                      # 箭头形状


def plot_node(node_txt, center_pt, parent_pt, node_type):
    create_plot.ax1.annotate(node_txt, xy=parent_pt, xycoords='axes fraction',
                             xytext=center_pt, textcoords='axes fraction',
                             va='center', ha='center', bbox=node_type, arrowprops=arrow_args)


def get_num_leafs(my_tree):
    num_leafs = 0
    # print(num_leafs, '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    first_str = list(my_tree.keys())[0]
    second_dict = my_tree[first_str]
    for key in second_dict.keys():
        if type(second_dict[key]).__name__ == 'dict':
            num_leafs += get_num_leafs(second_dict[key])
        else:
            num_leafs += 1
        # print(num_leafs)
    return num_leafs


def get_tree_depth(my_tree):
    max_depth = 0
    first_str = list(my_tree.keys())[0]
    second_dict = my_tree[first_str]
    for key in second_dict.keys():
        if type(second_dict[key]).__name__ == 'dict':
            this_depth = 1 + get_tree_depth(second_dict[key])
        else:
            this_depth = 1
        if this_depth > max_depth:
            max_depth = this_depth
    return max_depth


def retrieve_tree(i):
    list_trees = [{'no surfacing': {0: 'no', '1': {'flippers': {0: 'no', 1: 'yes'}}}},
                  {'no surfacing': {0: 'no', '1': {'flippers': {0: {'head': {0: 'no', 1: 'yes'}}, 1: 'no'}}}}]
    return list_trees[i]


def plot_mid_text(cntr_pt, parent_pt, txt_string):
    x_mid = (parent_pt[0]-cntr_pt[0])/2.0 + cntr_pt[0]
    y_mid = (parent_pt[1]-cntr_pt[1])/2.0 + cntr_pt[1]
    create_plot.ax1.text(x_mid, y_mid, txt_string)


def plot_tree(my_tree, parent_pt, node_txt):
    num_leafs = get_num_leafs(my_tree)
    depth = get_tree_depth(my_tree)
    first_str = list(my_tree.keys())[0]
    cntr_pt = plot_tree.x_off + (1.0 + float(num_leafs))/2.0/plot_tree.total_w, plot_tree.y_off
    plot_mid_text(cntr_pt, parent_pt, node_txt)
    plot_node(first_str, cntr_pt, parent_pt, decision_node)
    second_dict = my_tree[first_str]
    plot_tree.y_off = plot_tree.y_off - 1.0/plot_tree.total_d
    for key in second_dict.keys():
        if type(second_dict[key]).__name__ == 'dict':
            plot_tree(second_dict[key], cntr_pt, str(key))
        else:
            plot_tree.x_off = plot_tree.x_off + 1.0/plot_tree.total_w
            plot_node(second_dict[key], (plot_tree.x_off, plot_tree.y_off), cntr_pt, leaf_node)
            plot_mid_text((plot_tree.x_off, plot_tree.y_off), cntr_pt, str(key))
        plot_tree.y_off = plot_tree.y_off + 1.0/plot_tree.total_d


def create_plot(in_tree):
    fig = plt.figure(1, facecolor='white')                                         # 定义一个画布，背景色白色
    fig.clf()                                                                      # 清除画布
    axprops = dict(xticks=[], yticks=[])                                           # 在寻找这个是干啥的
    # 类、函数的一个属性，全局可用，subplot定义子画布，1行1列中的第一个画布，
    create_plot.ax1 = plt.subplot(111, frameon=False, **axprops)        
    plot_tree.total_w = float(get_num_leafs(in_tree))                              # 全局属性的叶子个数
    plot_tree.total_d = float(get_tree_depth(in_tree))                             # 树的深度
    plot_tree.x_off = -0.5/plot_tree.total_w                                       # 可能是每个叶子的偏移度
    plot_tree.y_off = 1.0                                                          # 可能是树高，每次往下，需要相减
    # plot_node('决策节点', (0.5, 0.1), (0.1, 0.5), decision_node)
    # plot_node('叶节点', (0.8, 0.1), (0.3, 0.8), leaf_node)
    plot_tree(in_tree, (0.5, 1.0), '')
    plt.show()


my_trees = retrieve_tree(0)
create_plot(my_trees)



